from datetime import datetime
import json
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from typing import List, Dict, Any
import requests
import time

app = FastAPI()

app.mount("/static", StaticFiles(directory="static"), name="static")


@app.get("/", response_class=HTMLResponse)
async def read_index():
    with open("static/index.html") as f:
        return HTMLResponse(content=f.read(), status_code=200)

PROMPT = """
你是一位专业的 Prompt Engineer，用户将发送一篇包含多个句子或段落的文章。你需要根据文章内容生成多样化的问题，并根据上下文分别提供 AI 风格的答案，同时标记出答案依据的上下文内容。

具体要求如下：

- **问题生成**：
  - 确保生成的问题与文章内容相关，并且对应的答案能够在原文中找到。问题应当有长有短，像真实人类提出的问题。
  - 避免生成过于直白和简单的问题。答案可能基于一个句子，也可能基于多个句子。
  - 要求生成的问题必须要包含设备名称或具体的维修过程的主体，避免生成模糊问题（如果是维修过程需要指明维修的设备或者维修主体）。
  - 如果有关于维修的问题要明确具体是什么维修，不要出现类似："这次维修。。。"这样的问题。（要明确维修的主体与过程（具体维修什么内容，比如维修“电流不稳”故障时。。。））

- **答案生成**：
  - 答案应符合主流 AI Assistant 的风格，在适当的地方使用换行符和 Markdown 等格式以提高美观度和可读性。
  - 确保答案能够在原文中找到且不包含无关或错误内容，同时尽量详细。
  - 注意不要改变原文的本意，更不能捏造事实。
  - 期望生成的比较多的问题和答案是，针对某一设备名称，有什么样的故障描述。针对每一个设备具体的故障描述，能够回答出相应的维修记录和举措。如果这个维修记录是内修，则不需要计算维修费用。如果是外修，还需要回答出相应的计算费用，看看是如何花费的。

- **引用内容**：
  - 答案依据的上下文内容应为完整的句子，而非不完整的文本片段。这些句子可能来自文章的多个部分。
  - 每个引用内容一般不超过500字。

- **设备名称**：
  - 每个生成的问题必须包含设备名称，明确是针对哪个设备提出的问题。
  - 避免模糊问题，例如“第几次维修…”，需要明确具体设备和故障情况。

- **维修记录**：
  - 针对设备的具体故障描述，生成相应的维修记录和措施。如果是内修，则不需要计算维修费用。如果是外修，还需提供费用计算。
  
- **分析过程**：
  - 生成的问题和答案需要 step by step 地进行分析，避免跳跃步骤。

### 示例

假设用户提供了一篇关于某种工业设备故障及维修记录的文章，示例如下：

**输入文章**：
> 5807-9040二氧化碳焊机发生电气类故障，故障描述为《更换主板》，采用外修方式维修，维修后维修人填写的记录为《经查为焊机线路板短路故障，更换维修线路板后，恢复设备使用。》。总花费为0.00(人员花费) + 193.16(备件花费) = 193.16(总花费)，其中：更换备件有：主P板维修，费用分别为：1(更换备件数量) * 193.16(备件单价) = 193.16(备件花费)

**生成的问题和答案**：

**问题1**：5807-9040二氧化碳焊机发生了什么类型的故障？  
**答案**：5807-9040二氧化碳焊机发生了电气类故障。  

**问题2**：5807-9040二氧化碳焊机的故障具体是什么？  
**答案**：焊机线路板短路故障。  

**问题3**：5807-9040二氧化碳焊机的维修方式是什么？  
**答案**：采用外修方式维修。  

**问题4**：维修人员在维修后记录了什么？  
**答案**：维修人员记录是《经查为焊机线路板短路故障，更换维修线路板后，恢复设备使用》。  

**问题5**：5807-9040二氧化碳焊机的总维修费用是多少？  
**答案**：总维修费用为193.16元

**问题6**：5807-9040二氧化碳焊机更换了哪些备件，费用是多少？  
**答案**：更换了主P板维修，费用为193.16元。  

通过这种方式，你可以生成与原文内容紧密相关的问题和答案，以提高问题生成的多样性和答案的准确性。
"""


def process_text_with_model(text: str, model_name: str) -> List[Dict[str, Any]]:
    url = "http://1.15.125.13:8012/v1/chat/completions"
    headers = {"Content-Type": "application/json"}
    if model_name == 'chatglm-4-9b-chat':
        headers["Authorization"] = "Bearer sk-3033&5004"
    payload = {
        "model": model_name,
        "messages": [
            {
                "role": "user",
                "content": f"{PROMPT}\n\n{text}",
                "name": "user"
            }
        ],
        "temperature": 0.8,
        "top_p": 0.8,
        "max_tokens": 0,
        "stream": False,
        "tools": {},
        "tool_choice": "None",
        "repetition_penalty": 1.1
    }
    response = requests.post(url, json=payload, headers=headers)
    return response.json()


@app.post("/uploadfiles/")
async def create_upload_files(files: List[UploadFile] = File(...),
                              model_name: str = "chatglm-4-9b-chat"):
    start_time = time.time()  # 记录开始时间
    qa_pairs = []
    for file in files:
        contents = await file.read()
        text = contents.decode("utf-8")
        segments = text.split("\n\n")  # 简单分段
        for segment in segments:
            qa_results = process_text_with_model(segment, model_name)
            for choice in qa_results['choices']:
                message_content = choice['message']['content']
                # 假设大模型返回的是分段后的 Q&A 对，使用 '**' 进行提取
                questions_and_answers = message_content.split('\n\n')
                for qa in questions_and_answers:
                    if '**问题' in qa and '**答案**' in qa:
                        question = qa.split('**问题')[1].split('**')[1].strip()
                        answer = qa.split('**答案**')[1].strip()
                        # 移除开头的冒号和多余的空格
                        if question.startswith('：'):
                            question = question[1:].strip()
                        if answer.startswith('：'):
                            answer = answer[1:].strip()
                        # 移除引用内容和多余的说明
                        answer = answer.split("依据：")[0].split("参考上文：")[0].strip()
                        qa_pairs.append({
                            "instruction": question,
                            "input": "",
                            "output": answer
                        })
    # 保存为 JSON 文件
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    output_file = f"qa_pairs_{timestamp}.json"
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(qa_pairs, f, ensure_ascii=False, indent=4)

    end_time = time.time()  # 记录结束时间
    elapsed_time = end_time - start_time  # 计算总用时

    return {"message": f"QA pairs saved to {output_file}", "qa_pairs": qa_pairs, "elapsed_time": elapsed_time}

# 添加 POST 请求以测试固定内容的响应


@app.post("/test-chat/")
async def test_chat(model_name: str = "chatglm-4-9b-chat"):
    url = "http://1.15.125.13:8012/v1/chat/completions"
    headers = {"Content-Type": "application/json"}
    if model_name == 'chatglm-4-9b-chat':
        headers["Authorization"] = "Bearer sk-3033&5004"
    payload = {
        "model": model_name,
        "messages": [
            {
                "role": "user",
                "content": "你好",
                "name": "string"
            }
        ],
        "temperature": 0.8,
        "top_p": 0.8,
        "max_tokens": 0,
        "stream": False,
        "tools": {},
        "tool_choice": "None",
        "repetition_penalty": 1.1
    }
    response = requests.post(url, json=payload, headers=headers)
    return response.json()


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8080)
